import pickle
import numpy as np
import json

from numpy.core.fromnumeric import shape

def gen_poseC3d_annotation(dataset):
    all_keypoints = None
    all_scores = None
    for data in dataset:
        one_keypoints = data['keypoints'].numpy()
        list_one_keypoints = []
        list_one_keypoints.append(one_keypoints)
        one_scores = data['kp_score'].numpy().reshape(17)
        
        #list_one_keypoints: 17, 2
        #len(list_score): 17
        one_keypoints = np.array(list_one_keypoints)

        print(one_keypoints.shape)

        if all_keypoints is None:
            all_keypoints = one_keypoints
        else:
            all_keypoints = np.concatenate((all_keypoints, one_keypoints), axis=0)

        one_scores = one_scores.reshape((1,) + one_scores.shape)
        if all_scores is None:
            all_scores = one_scores
        else:
            all_scores = np.concatenate((all_scores, one_scores), axis=0)

    all_keypoints = all_keypoints.reshape((1,) + all_keypoints.shape)
    all_scores = all_scores.reshape((1,) + all_scores.shape)

    d = {}
    d['keypoint'] = all_keypoints
    d['keypoint_score'] = all_scores
    d['total_frames'] = all_keypoints.shape[1]
    d['original_shape'] = (1280, 720)
    d['img_shape'] = (1280, 1080)
    d['label'] = 1
    d['frame_dir'] = '/home/primer/videos/action.mp4'

    with open ('/home/primer/code/action/action_recognizition/mmaction2/data/posec3d/one_anno.pkl', 'wb') as f:
        pickle.dump([d], f)

def offline_gen_poseC3d_annotation(dataset):
    all_keypoints = None
    all_scores = None
    for data in dataset:
        one_keypoints = np.array(data['keypoints'])
        list_one_keypoints = []
        list_one_keypoints.append(one_keypoints)
        one_scores = np.array(data['kp_score']).reshape(17)
        
        #list_one_keypoints: 17, 2
        #len(list_score): 17
        one_keypoints = np.array(list_one_keypoints)

        print(one_keypoints.shape)

        if all_keypoints is None:
            all_keypoints = one_keypoints
        else:
            all_keypoints = np.concatenate((all_keypoints, one_keypoints), axis=0)

        one_scores = one_scores.reshape((1,) + one_scores.shape)
        if all_scores is None:
            all_scores = one_scores
        else:
            all_scores = np.concatenate((all_scores, one_scores), axis=0)

    all_keypoints = all_keypoints.reshape((1,) + all_keypoints.shape)
    all_scores = all_scores.reshape((1,) + all_scores.shape)

    d = {}
    d['keypoint'] = all_keypoints
    d['keypoint_score'] = all_scores
    d['total_frames'] = all_keypoints.shape[1]
    d['original_shape'] = (1280, 720)
    d['img_shape'] = (1280, 1080)
    d['label'] = 1
    d['frame_dir'] = '/home/primer/videos/action.mp4'

    with open ('/home/primer/code/action/action_recognizition/mmaction2/data/posec3d/one_anno.pkl', 'wb') as f:
        pickle.dump([d], f)


def offline(address = '/home/primer/code/AlphaPose_bak/examples/468/alphapose-results.json'):
    with open(address) as f:
        dataset = json.load(f)

    all_keypoints = None
    all_scores = None
    for data in dataset:
        one_keypoints = data['keypoints']
        list_one_keypoints = []
        list_score = []
        for i in range(0, len(one_keypoints), 3):
            x, y = one_keypoints[i], one_keypoints[i+1]
            score = one_keypoints[i+2]
            list_one_keypoints.append([x, y])
            list_score.append(score)
        
        #list_one_keypoints: 17, 2
        #len(list_score): 17
        one_keypoints = np.array(list_one_keypoints)
        one_keypoints = one_keypoints.reshape((1,) + one_keypoints.shape)

        if all_keypoints is None:
            all_keypoints = one_keypoints
        else:
            all_keypoints = np.concatenate((all_keypoints,one_keypoints), axis=0)

        one_scores = np.array(list_score)
        one_scores = one_scores.reshape((1,) + one_scores.shape)
        if all_scores is None:
            all_scores = one_scores
        else:
            all_scores = np.concatenate((all_scores, one_scores), axis=0)

    all_keypoints = all_keypoints.reshape((1,) + all_keypoints.shape)
    all_scores = all_scores.reshape((1,) + all_scores.shape)

    d = {}
    d['keypoint'] = all_keypoints
    d['keypoint_score'] = all_scores
    d['total_frames'] = all_keypoints.shape[1]
    d['original_shape'] = (1280, 720)
    d['img_shape'] = (1280, 1080)
    d['label'] = 1
    d['frame_dir'] = '/home/primer/videos/action.mp4'

    with open ('mmaction2/one_anno_test/one_annotation.pkl','wb') as f:
        pickle.dump([d], f)


if __name__ == '__main__':
    offline()